import serial
import os
import time
import struct
import sys

class YModemSender:
    """优化ACK检测的YModem发送端实现"""
    # 协议控制字符
    SOH = 0x01
    STX = 0x02
    EOT = 0x04
    ACK = 0x06
    NAK = 0x15
    CAN = 0x18
    C = 0x43

    def __init__(self, serial_port, baudrate=115200, timeout=5):
        self.serial_port = serial_port
        self.baudrate = baudrate
        self.timeout = timeout  # 延长超时时间，确保能收到慢响应
        self.packet_size = 1024
        self.transfer_aborted = False
        self.serial = None

    def open(self):
        """打开串口并配置优化参数"""
        if self.serial is None or not self.serial.is_open:
            self.serial = serial.Serial(
                port=self.serial_port,
                baudrate=self.baudrate,
                timeout=self.timeout,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                bytesize=serial.EIGHTBITS,
                # 关闭流控以减少延迟和冲突
                xonxoff=False,
                rtscts=False,
                dsrdtr=False
            )
            # 清空缓冲区但保留足够时间
            time.sleep(0.1)
            self.serial.flushInput()
            self.serial.flushOutput()
            print(f"串口 {self.serial_port} 已打开")
        return self.serial.is_open

    def close(self):
        """关闭串口"""
        if self.serial and self.serial.is_open:
            self.serial.close()
            print(f"串口 {self.serial_port} 已关闭")
        self.serial = None

    def _crc16(self, data):
        """计算CRC16校验和"""
        crc = 0x0000
        for byte in data:
            crc ^= (byte << 8)
            for _ in range(8):
                if crc & 0x8000:
                    crc = (crc << 1) ^ 0x1021
                else:
                    crc <<= 1
                crc &= 0xFFFF
        return crc

    def _read_until_ack(self, max_attempts=5):
        """
        优化的ACK读取逻辑：
        1. 忽略无关数据，只关注ACK
        2. 多次尝试，应对数据延迟到达
        3. 显示原始接收数据，便于调试
        """
        if not self.serial or not self.serial.is_open:
            print("错误：串口未打开")
            return False

        for attempt in range(max_attempts):
            start_time = time.time()
            buffer = b''
            
            # 循环读取直到超时或找到ACK
            while time.time() - start_time < self.timeout:
                # 读取所有可用数据（关键改进）
                if self.serial.in_waiting > 0:
                    buffer += self.serial.read(self.serial.in_waiting)
                    
                    # 检查缓冲区中是否包含ACK（即使有其他数据）
                    if self.ACK in buffer:
                        # print(f"检测到ACK（尝试{attempt+1}/{max_attempts}）")
                        # # 打印原始数据用于调试（可在实际使用时注释）
                        # print(f"原始响应数据: {[hex(b) for b in buffer]}")
                        return True
                
                # 短暂延迟，减少CPU占用
                time.sleep(0.01)
            
            # 未找到ACK，显示尝试结果
            print(f"未检测到ACK（尝试{attempt+1}/{max_attempts}）")
            if buffer:
                print(f"收到无关数据: {[hex(b) for b in buffer]}")
            
            # 清空缓冲区，准备下一次尝试
            self.serial.flushInput()
        
        return False

    def _send_packet(self, sequence, data):
        """发送数据包并优化ACK检测"""
        if not self.serial or not self.serial.is_open:
            print("错误：串口未打开")
            return False

        # 确定数据包类型
        if len(data) == 128:
            packet_type = self.SOH
        elif len(data) == 1024:
            packet_type = self.STX
        else:
            return False

        # 计算CRC
        crc = self._crc16(data)
        
        # 构建数据包
        packet = bytearray()
        packet.append(packet_type)
        packet.append(sequence & 0xFF)
        packet.append((~sequence) & 0xFF)
        packet.extend(data)
        packet.append((crc >> 8) & 0xFF)
        packet.append(crc & 0xFF)

        # 发送数据包前确保缓冲区干净
        self.serial.flushInput()
        self.serial.flushOutput()
        
        # 发送数据包
        self.serial.write(packet)
        self.serial.flush()
        print(f"已发送数据包 {sequence} (大小: {len(data)}字节)")

        # 等待ACK（使用优化的检测方法）
        return self._read_until_ack()

    def send_file(self, filename):
        """发送文件，重点优化文件信息包后的ACK检测"""
        if not self.serial or not self.serial.is_open:
            print("错误：串口未打开，无法发送文件")
            return False

        if not os.path.exists(filename):
            print(f"错误: 文件 {filename} 不存在")
            return False

        try:
            file_size = os.path.getsize(filename)
            file_name = os.path.basename(filename)
            print(f"准备发送文件: {file_name}, 大小: {file_size}字节")

            # 步骤1: 等待接收方的初始同步信号（C）
            print("等待接收方同步信号...")
            # 简化同步逻辑，专注于文件信息包的ACK问题
            if not self._wait_for_c_sync():
                return False

            # 步骤2: 发送文件信息包（序号0）- 重点优化点
            print("发送文件信息包...")
            info_str = f"{file_name}\0{file_size}\0"
            info_data = info_str.encode('ascii')
            if len(info_data) < 128:
                info_data += b'\x00' * (128 - len(info_data))
            
            # 对文件信息包使用更多的重试次数
            info_packet_success = False
            for retry in range(3):  # 最多重试3次
                if self._send_packet(0, info_data):
                    info_packet_success = True
                    break
                print(f"文件信息包发送失败，重试 {retry+1}/3")
                time.sleep(1)  # 重试前等待设备准备
            
            if not info_packet_success:
                print("文件信息包多次发送失败，终止传输")
                return False

            # # 后续步骤保持不变
            # if not self._wait_for_c_sync():
            #     return False

            print("开始发送文件数据...")
            with open(filename, 'rb') as f:
                sequence = 1
                total_packets = (file_size + self.packet_size - 1) // self.packet_size
                
                while True:
                    if self.transfer_aborted:
                        return False
                        
                    data = f.read(self.packet_size)
                    if not data:
                        break
                        
                    if len(data) < self.packet_size:
                        data += b'\x1A' * (self.packet_size - len(data))
                        
                    if not self._send_packet(sequence, data):
                        print(f"数据包 {sequence} 发送失败")
                        return False
                        
                    progress = (sequence / total_packets) * 100
                    print(f"传输进度: {progress:.1f}% ({sequence}/{total_packets})", end='\r')
                    sequence += 1

            print("\n所有数据数据包发送完成")
                
            # 发送EOT和结束包的逻辑保持不变
            if not self._send_eot():
                return False

            print("发送结束包...")
            end_data = b'\x00' * 128
            if not self._send_packet(0, end_data):
                print("结束包发送失败")
                return False
            
            print("文件发送成功!")
            return True

        except Exception as e:
            print(f"传输过程中发生错误: {str(e)}")
            return False

    def _wait_for_c_sync(self):
        """等待接收方的C同步信号"""
        for _ in range(5):
            if self.serial.in_waiting > 0:
                data = self.serial.read(self.serial.in_waiting)
                if self.C in data:
                    print("收到同步信号'C'")
                    return True
                elif self.CAN in data:
                    self.transfer_aborted = True
                    print("收到取消信号")
                    return False
            
            # self.serial.write(bytes([self.C]))
            # self.serial.flush()
            time.sleep(1)
        
        print("未收到同步信号'C'")
        return False

    def _send_eot(self):
        """发送EOT信号"""
        print("发送传输结束标志...")
        # self.serial.write(bytes([self.EOT]))
        # self.serial.flush()
        # response = self.serial.read(1)
        # if response != bytes([self.NAK]):
        #     print("第一次发送EOT未收到NAK")
        #     return False
            
        self.serial.write(bytes([self.EOT]))
        self.serial.flush()
        response = self.serial.read(1)
        if response != bytes([self.ACK]):
            print("第二次发送EOT未收到ACK")
            return False
            
        return True

class DeviceUpgrader:
    """设备升级管理器"""
    def __init__(self, serial_port, baudrate=115200):
        self.serial_port = serial_port
        self.baudrate = baudrate
        self.device_commands = {
            "mcu": "$mcu_ota,$end",
            "fpga": "$fpga_ota,$end",
            # "mcu": "esp32ota",
            # "fpga": "download"            
        }
        self.ota_ack = {
            "mcu": "MCU_OTA_READY",
            "fpga": "FPGA_OTA_READY"
        }
        self.ymodem = YModemSender(serial_port, baudrate)

    def add_device_command(self, device_type, command, ack_response=None):
        self.device_commands[device_type] = command
        if ack_response:
            self.ota_ack[device_type] = ack_response

    def enter_download_mode(self, device_type):
        if device_type not in self.device_commands:
            print(f"不支持的设备类型: {device_type}")
            return False

        if not self.ymodem.open():
            print("无法打开串口")
            return False

        try:
            cmd = self.device_commands[device_type]
            print(f"发送{device_type}进入OTA模式命令: {cmd}")
            
            self.ymodem.serial.flushInput()
            self.ymodem.serial.flushOutput()
            self.ymodem.serial.write(cmd.encode('utf-8'))
            self.ymodem.serial.flush()
            
            if device_type in self.ota_ack and self.ota_ack[device_type]:
                print(f"等待{device_type}确认响应...")
                response = b''
                start_time = time.time()
                while time.time() - start_time < 5:
                    if self.ymodem.serial.in_waiting > 0:
                        response += self.ymodem.serial.read(self.ymodem.serial.in_waiting)
                        if self.ota_ack[device_type].encode() in response:
                            print(f"{device_type}已进入OTA模式")
                            return True
                    time.sleep(0.1)
                print(f"{device_type}响应: {response.decode(errors='ignore')}")            
            return True
            
        except Exception as e:
            print(f"发送OTA命令失败: {str(e)}")
            return False

    def upgrade(self, device_type, firmware_file, packet_size=1024):
        try:
            if not self.ymodem.open():
                return False
                
            if not self.enter_download_mode(device_type):
                return False
                
            self.ymodem.packet_size = packet_size
            return self.ymodem.send_file(firmware_file)
            
        finally:
            time.sleep(1)
            self.ymodem.close()

if __name__ == "__main__":
    if len(sys.argv) != 5:
        print(f"用法: {sys.argv[0]} <串口> <波特率> <设备类型> <固件文件>")
        print(f"示例: {sys.argv[0]} COM3 115200 mcu  RFcarrier*.bin")
        print(f"示例: {sys.argv[0]} COM3 115200 fpga uav_top*.bin")
        sys.exit(1)

    serial_port = sys.argv[1]
    baudrate = int(sys.argv[2])
    device_type = sys.argv[3]
    firmware_file = sys.argv[4]

    upgrader = DeviceUpgrader(serial_port, baudrate)
    success = upgrader.upgrade(device_type, firmware_file)
    sys.exit(0 if success else 1)